package com.eian.boot.web.config.thread;

import com.alibaba.ttl.threadpool.TtlExecutors;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.boot.web.embedded.tomcat.TomcatProtocolHandlerCustomizer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.annotation.AsyncConfigurer;
import org.springframework.scheduling.annotation.EnableAsync;

import java.util.Map;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;

@Slf4j
@Configuration
@EnableAsync(proxyTargetClass = true)
public class AsyncConfig implements AsyncConfigurer {
    /**
     * @author alex.meng
     * @description @Async 使用虚拟线程
     */
    @Override
    @Primary
    @Bean("taskExecutor")
    public Executor getAsyncExecutor() {
        ThreadFactory factory = Thread.ofVirtual()
                .name("task-handle-", 0)
                .factory();
        ExecutorService virtualThreadExecutor = Executors.newThreadPerTaskExecutor(factory);
        log.info("异步任务(@Async) 虚拟线程池初始化完成");
        return TtlExecutors.getTtlExecutorService(virtualThreadExecutor);
    }

    @Bean
    public TomcatProtocolHandlerCustomizer<?> protocolHandlerVirtualThreadExecutorCustomizer() {
        log.info("SpringBoot WEB 虚拟线程已开启");
        return protocolHandler -> protocolHandler.setExecutor(Executors.newVirtualThreadPerTaskExecutor());
    }

    @Bean("mdcTaskDecorator")
    public TaskDecorator mdcTaskDecorator() {
        return runnable -> {
            // 捕获提交任务时的MDC上下文
            Map<String, String> contextMap = MDC.getCopyOfContextMap();
            return () -> {
                try {
                    // 任务执行前恢复上下文
                    if (contextMap != null) {
                        MDC.setContextMap(contextMap);
                    }
                    runnable.run();
                } finally {
                    // 任务执行后清理上下文
                    MDC.clear();
                }
            };
        };
    }
}
